from collections import ChainMap
import datetime
import joblib
import numpy as np
from random import random
import torch
from torch import tensor
from tjdutils.utils import current_time, check_path


class Saver(object):
    def __init__(self, key1, value1, key2, value2, model, trainer, loader, output, constructor_path, deep_learning: bool = True):
        self.key = ''.join([key1, key2])
        self.path = ''.join([constructor_path,'/model_', current_time(date_format = "us"),'_',str(random())[2:8]])
        if deep_learning:

            model_path = check_path(''.join([self.path, '.pth']))
            torch.save(model.state_dict(), model_path)
            y_test_predictions = trainer.test(loader.x_test, model_path)
            y_test_predictions = self.to_array(y_test_predictions)
            index = ['y_train_predictions', 'y_train_targets', 'y_valid_predictions', 'y_valid_targets',
                     'y_test_targets', 'train_loss', 'valid_loss']
            output = self.to_array(output)
            model_output = dict(zip(index, output))
            model_output['y_test_predictions'] = y_test_predictions
            #model_output['selector'] = value1
            model_output['model_path'] = model_path
        else:   

            model_path = check_path(''.join([self.path, '.joblib']))
            joblib.dump(model, model_path)
            index = ['y_train_predictions', 'y_train_targets', 'y_valid_predictions', 'y_valid_targets',
                     'y_test_predictions', 'y_test_targets', 'train_loss', 'valid_loss']
            model_output = dict(zip(index, output))
            #model_output['selector'] = value1
            model_output['model_path'] = model_path

        #vv = ChainMap(value1, value2, model_output)
        vv = {**value1, **value2, **model_output}
        vv["trainer_key"] = ''.join([key1, key2])
        self.trainer_space = vv

    @staticmethod
    def to_array(x):
        if isinstance(x, tuple):
            x = (np.array(i).reshape(-1, 1) for i in x) if (np.array(i).ndim == 0 or 1 for i in x) else (np.array(i) for i in x)
        else:
            x = np.array(x).reshape(-1, 1) if np.array(x).ndim == 0 or 1 else np.array(x)
        return x







# 1. 先建立一个字典，保存三个参数：
#
# state = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
#
# 2.调用torch.save():
#
# torch.save(state, dir)
#
# 其中dir表示保存文件的绝对路径+保存文件名，如'/home/qinying/Desktop/modelpara.pth'
#
# 二、
#
# 当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。
#
# checkpoint = torch.load(dir)
#
# model.load_state_dict(checkpoint['net'])
#
# optimizer.load_state_dict(checkpoint['optimizer'])
#
# start_epoch = checkpoint['epoch'] + 1


